import torch

from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
from torch import nn, optim
#from lenet5 import Lenet5
from resnet import ResNet18

def main():
    batch_size = 32
    # 这个是一次加载一张的数据，所以后续还要定义怎么按批次加载
    cifar_train = datasets.CIFAR10('cifar', True, transform=transforms.Compose([
        transforms.Resize((32, 32)), # 转换数据格式
        transforms.ToTensor(), # ToTensor，Tensor为计算的基本单元
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ) # 标准化，参数都是设定好的
    ]), download=True)
    # 数据加载器，# shuffle设置随机化
    cifar_train = DataLoader(cifar_train, batch_size=batch_size, shuffle=True)

    cifar_test = datasets.CIFAR10('cifar', False, transform=transforms.Compose([
        transforms.Resize((32, 32)),  # 转换数据格式
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ) # 标准化
    ]), download=True)
    cifar_test = DataLoader(cifar_test, batch_size=batch_size, shuffle=True)  # shuffle设置随机化

    x, label = iter(cifar_train).next()
    print("x: ", x.shape, "  label: ", label.shape)

    device = torch.device('cuda') # 命名一个GPU设备
    #model = Lenet5().to(device)
    model = ResNet18().to(device) # 将模型导入到GPU
    criteon = nn.CrossEntropyLoss().to(device)
    print(model)
    optimizer = optim.Adam(model.parameters(), lr=1e-3) # 定义优化器

    for epoch in range(2):
        model.train() # 将模型变为训练模式
        for batch_index, (x, label) in enumerate(cifar_train):
            # [b, 3, 32, 32], [b]
            x, label = x.to(device), label.to(device)

            logits = model(x)
            # logits:   [b, 10]
            # label:    [b]
            # loss:     tensor scalar
            loss = criteon(logits, label)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(epoch, loss.item())

        model.eval() # 将模型变为测试模式
        with torch.no_grad(): # 告诉torch不需要计算梯度
            # test
            total_correct = 0
            total_num = 0
            for x, label in cifar_test:
                # [b, 3, 32, 32], [b]
                x, label = x.to(device), label.to(device)

                # [b, 10]
                logits = model(x)
                # [b]
                pred = logits.argmax(dim=1)
                # [b] vs [b] => scalar tensor
                total_correct += torch.eq(pred, label).float().sum().item()
                total_num += x.size(0)

            acc = total_correct / total_num
            print("test", epoch, acc)

if __name__ == '__main__':
    main()
